import tensorflow as tf
import tensorboard as tb
import numpy as np


##################################  ACTOR  #################################################
class Actor(object):
    def __init__(self,S,S_,sess, action_dim, action_bound, learning_rate, replacement):
        self.sess = sess    # Tensorflow Session
        self.a_dim = action_dim     # 动作维度[角度2维连续，攻击1维onehot]
        self.action_bound = action_bound    # 动作输出的幅度
        self.lr = learning_rate     # 学习率
        self.replacement = replacement      # 两个网络的融合模式
        self.t_replace_counter = 0      # 替换回合计时
        self.S_input,self.S_input_=S,S_
        self.a_eval,self.a_target=self.inference()

        self.e_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='Actor/eval_net')
        self.t_params = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='Actor/target_net')

        if self.replacement['name'] == 'hard':
            self.t_replace_counter = 0
            self.hard_replace = [tf.assign(t, e) for t, e in zip(self.t_params, self.e_params)]
        else:
            self.soft_replace = [tf.assign(t, (1 - self.replacement['tau']) * t + self.replacement['tau'] * e)
                                 for t, e in zip(self.t_params, self.e_params)]

    def inference(self):
        with tf.variable_scope('Actor'):
            # input s, output a
            a_eval = self._build_net(self.S_input, scope='eval_net', trainable=True)
            # input s_, output a_, get a_ for critic
            a_target = self._build_net(self.S_input_, scope='target_net', trainable=False)
        return a_eval,a_target

    def _build_net(self, s, scope, trainable):
        with tf.variable_scope(scope):
            init_w = tf.random_normal_initializer(0., 0.3)
            init_b = tf.constant_initializer(0.1)
            l1 = tf.layers.dense(s, 256, activation=tf.nn.relu,
                                  kernel_initializer=init_w, bias_initializer=init_b, name='l1',
                                  trainable=trainable)
            l2=tf.layers.dense(l1, 256, activation=tf.nn.relu,
                                  kernel_initializer=init_w, bias_initializer=init_b, name='l2',
                                  trainable=trainable)
            with tf.variable_scope('a_'):
                ang_action = tf.layers.dense(l2, self.a_dim[0], activation=tf.nn.tanh, kernel_initializer=init_w,
                                          bias_initializer=init_b, name='a', trainable=trainable)
                bit_action=tf.layers.dense(l2, 2, activation=tf.nn.tanh, kernel_initializer=init_w,
                                          bias_initializer=init_b, name='a', trainable=trainable)   # 输出两维的[0,1]    # TODO 添加其它onehot输出需要修改
                scaled_a = tf.multiply(ang_action, self.action_bound, name='scaled_a')

        return scaled_a,bit_action

    def choose_action(self, state):
        s = state[np.newaxis, :]  # 一维转二维，才能输入tf
        return self.sess.run(self.a_eval, feed_dict={self.S_input: state})  # tf的输出也是二维的，需要再行转化

    def add_grad_to_graph(self, a_grads):
        """
        a_grads来自Critic网络
        :param a_grads:
        :return:
        """
        with tf.variable_scope('policy_grads'):
            # ys = policy;
            # xs = policy's parameters;
            # self.a_grads = 可以得到更大Q值的梯度方向
            # tf.gradients 计算 dys/dxs with a initial gradients for ys, so this is dq/da * da/dparams
            self.policy_grads = tf.gradients(ys=self.a_eval, xs=self.e_params, grad_ys=a_grads)

        with tf.variable_scope('A_train'):
            opt = tf.train.AdamOptimizer(-self.lr)  # (- learning rate) for ascent policy
            self.train_op = opt.apply_gradients(zip(self.policy_grads, self.e_params))

    def learn(self, state):  # batch update
        self.sess.run(self.train_op, feed_dict={self.S_input: state})
        if self.replacement['name'] == 'soft':
            self.sess.run(self.soft_replace)
        else:
            if self.t_replace_counter % self.replacement['rep_iter_a'] == 0:
                self.sess.run(self.hard_replace)
            self.t_replace_counter += 1

######################################  CRITIC  ############################################
class Critic(object):
    def __init__(self,S,S_,sess, state_dim, action_dim, learning_rate, gamma, replacement, a, a_):
        self.sess = sess
        self.s_dim = state_dim
        self.a_dim = action_dim
        self.lr = learning_rate
        self.gamma = gamma
        self.replacement = replacement
        self.S_input, self.S_input_ = S, S_

        self.c_eval,self.c_target=self.inference(a,a_)

    def inference(self,a,a_):
        with tf.variable_scope('Critic'):
            # Input (s, a), output q
            a = a
            c_eval = self._build_net(self.S_input, a, 'eval_net', trainable=True)

            # Input (s_, a_), output q_ for q_target
            c_target = self._build_net(self.S_input_, a_, 'target_net', trainable=False)    # c_target的a来自于a_target

        return  c_eval,c_target

    def _build_net(self, s, a, scope, trainable):
        with tf.variable_scope(scope):
            init_w = tf.random_normal_initializer(0., 0.1)
            init_b = tf.constant_initializer(0.1)

            with tf.variable_scope('l1'):
                n_l1=256
                l1=tf.layers.dense(s,n_l1,activation=tf.nn.relu,
                                   kernel_initializer=init_w,bias_initializer=init_b,name='l1',trainable=trainable)

            with tf.variable_scope('l2'):
                n_l2 = 256
                w2_s = tf.get_variable('w1_s', [n_l1, n_l2], initializer=init_w, trainable=trainable)
                w2_a = tf.get_variable('w1_a', [self.a_dim, n_l2], initializer=init_w, trainable=trainable)
                b1 = tf.get_variable('b1', [1, n_l1], initializer=init_b, trainable=trainable)
                net = tf.nn.relu(tf.matmul(l1, w2_s) + tf.matmul(a, w2_a) + b1)

            with tf.variable_scope('q'):
                q = tf.layers.dense(net, 1, kernel_initializer=init_w, bias_initializer=init_b, trainable=trainable)   # Q(s,a)
        return q


    def learn(self, s, a, r, s_):
        self.sess.run(self.train_op, feed_dict={S: s, self.a: a, R: r, S_: s_})
        if self.replacement['name'] == 'soft':
            self.sess.run(self.soft_replacement)
        else:
            if self.t_replace_counter % self.replacement['rep_iter_c'] == 0:
                self.sess.run(self.hard_replacement)
            self.t_replace_counter += 1